from constants import *
from codes import *


class CodeWriter():
    def __init__(self, file) -> None:
        self.file = open(file, 'w')
        self.label_index = 0

    def writeArithmetic(self, cmd):
        if cmd in ['add', 'sub', 'and', 'or']:
            cmd = two_operand(cmd)
        elif cmd in ['neg', 'not']:
            cmd = one_operand(cmd)
        elif cmd in ['eq', 'lt', 'gt']:
            cmd = jmp(cmd, get_new_label(self.label_index))
            self.label_index = self.label_index + 1
        self.file.write(cmd)

    def writePushPop(self, srcfile, cmd_type, segment, offset):
        if cmd_type == C_PUSH:
            self.file.write(push(srcfile, segment, offset))
        elif cmd_type == C_POP:
            self.file.write(pop(srcfile, segment, offset))

    def Close(self):
        self.file.close()
